package database

import (
	"database/sql"
	"fmt"
	"strings"

	_ "github.com/go-sql-driver/mysql" // MySQL驱动
)

// DB 是数据库连接的包装
type DB struct {
	*sql.DB
	driverName string
	dbName     string
}

// Connect 连接到数据库
func Connect(connStr string) (*DB, error) {
	// 默认使用MySQL
	driverName := "mysql"
	
	// 从连接字符串中提取数据库名
	dbName := extractDBName(connStr)
	if dbName == "" {
		return nil, fmt.Errorf("无法从连接字符串中提取数据库名")
	}

	db, err := sql.Open(driverName, connStr)
	if err != nil {
		return nil, err
	}

	if err := db.Ping(); err != nil {
		db.Close()
		return nil, err
	}

	return &DB{
		DB:         db,
		driverName: driverName,
		dbName:     dbName,
	}, nil
}

// GetTables 获取数据库中的所有表
func (db *DB) GetTables() ([]string, error) {
	query := "SHOW TABLES"
	rows, err := db.Query(query)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var tables []string
	for rows.Next() {
		var table string
		if err := rows.Scan(&table); err != nil {
			return nil, err
		}
		tables = append(tables, table)
	}

	return tables, nil
}

// GetTableStructure 获取表的结构信息
func (db *DB) GetTableStructure(tableName string) ([]Column, error) {
	query := fmt.Sprintf("DESCRIBE `%s`", tableName)
	rows, err := db.Query(query)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var columns []Column
	for rows.Next() {
		var col Column
		var null, key, extra, defaultVal sql.NullString
		if err := rows.Scan(&col.Field, &col.Type, &null, &key, &defaultVal, &extra); err != nil {
			return nil, err
		}
		
		col.Null = null.String
		col.Key = key.String
		col.Default = defaultVal.String
		col.Extra = extra.String
		
		columns = append(columns, col)
	}

	return columns, nil
}

// Column 表示表中的一列
type Column struct {
	Field   string
	Type    string
	Null    string
	Key     string
	Default string
	Extra   string
}

// 从连接字符串中提取数据库名
func extractDBName(connStr string) string {
	// 对于MySQL格式: user:password@tcp(host:port)/dbname
	parts := strings.Split(connStr, "/")
	if len(parts) > 1 {
		// 处理可能的参数
		dbNameParts := strings.Split(parts[len(parts)-1], "?")
		return dbNameParts[0]
	}
	return ""
}